/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.heron-language.com
*/

#ifndef LLSCHEME_REMOVE_RETURNS_HPP
#define LLSCHEME_REMOVE_RETURNS_HPP

namespace lls
{
	bool ContainsReturn(NonNull<SExpr> x) {
		if (x->Is<Return>()) {
			return true;
		}
		else {
			for (int i=0; i < x->GetNumChildren(); ++i) {
				SExpr* child = x->GetChild(i);
				if (!child->Is<Define>() && ContainsReturn(child)) 
					return true;
			}
			return false;
		}
	}

	void PlaceIfInTailPosition(NonNull<Seq> blk, int i) {
		Assert(blk->GetNumChildren() > i);
		Assert(blk->GetChild(i)->Is<If>());

		// already in tail position? 
		if (i == blk->GetNumChildren() - 1)
			return;

		If* ifs = blk->GetChild(i)->As<If>();
		NonNull<SExpr> trueBranch = ifs->TrueBranch();
		NonNull<SExpr> falseBranch = ifs->FalseBranch();
		
		// Move the rest of the statements in the code
		// block into the else branch
		
		// The assumption here is that all return statements
		// are in tail-position
		for (int j=i; j < blk->GetNumChildren(); ++j) {
			trueBranch->AddChild(blk->GetChild(j)->Clone());
		}
		for (int j=i; j < blk->GetNumChildren(); ++j) {
			falseBranch->AddChild(blk->GetChild(j)->Clone());
		}
		while (blk->GetNumChildren() > i) {
			blk->DeleteChild(i);
		}
	}

	void PlaceIfsInTailPosition(NonNull<SExpr> x) {
		x->ForEach(PlaceIfsInTailPosition);
		if (x->Is<Seq>()) {
			Seq* blk = x->As<Seq>();
			for (int i=0; i < blk->GetNumChildren() - 1; ++i) {
				if (blk->GetChild(i)->Is<If>()) {
					PlaceIfInTailPosition(blk, i);
				}
			}
		}
	}

	void CheckNoReturnsInLoops(NonNull<SExpr> x) {
		if (IsLoopConstruct(x)) {
			if (ContainsReturn(x)) {
				Error("Loops cannot contain return statements");
			}
		}
		else {
			x->ForEach(CheckNoReturnsInLoops);
		}
	}

	void RemoveStatementsAfterReturn(NonNull<SExpr> x) {
		// Assumption: all "if" statements are in tail position.
		if (x->Is<Seq>()) {
			Seq* blk = x->As<Seq>();
			int i=0;
			while (i < blk->GetNumChildren()) {
				if (blk->GetChild(i)->Is<Return>()) {
					while (i < blk->GetNumChildren()) {
						blk->DeleteChild(i);
					}
				}
				else {
					++i;				
				}
			}
		}
		else {
			x->ForEach(RemoveStatementsAfterReturn);
		}
	}

	void PlaceReturnsInTailPosition(NonNull<SExpr> x) {		
		// Make sure that no return statement occurs in a loop.
		// TODO: support loops with return statements
		CheckNoReturnsInLoops(x);

		// Recursively make sure that all return statements are in
		// tail position
		x->ForEach(PlaceReturnsInTailPosition);
		
		// first place any ifs in tail position
		// Note: this is inefficient, because it is a recursive algorithm so gets 
		// called moire often than needed. We don't care that much though.
		PlaceIfsInTailPosition(x);

		// The previous algorithm can leave superflous "return" statements.
		// Note: again another recursive / recursive algorithm.
		RemoveStatementsAfterReturn(x);
	}
}

#endif
